import subprocess
import os
import json


class AtcOm(object):
    def __init__(self):
        self.op_debug_level = 0

    def save_json(self, data, file_path):
        with open(file_path, "w") as f:
            json.dump(data, f, indent=4)

    def run_atc(self, params):
        format_str = "--%s=%s"
        param_list = []
        for key, value in params.items():
            param_str = format_str % (key, value)
            param_list.append(param_str)

        atc_param = " ".join(param_list)
        atc_cmd = "atc %s" % atc_param
        print("ATC CMD:", atc_cmd)
        return_code = subprocess.call(atc_cmd, shell=True, stderr=subprocess.PIPE, timeout=600)
        return return_code

    def generate_json(self, params, op_dir):
        json_name = params["op"].lower() + ".json"
        json_path = os.path.join(op_dir, json_name)
        self.save_json([params], json_path)
        return json_path

    def generate_om(self, params, op_dir, soc_version="Ascend910"):
        json_path = self.generate_json(params, op_dir)
        om_params = {"singleop": json_path, "output": op_dir, "soc_version": soc_version,
                     "precision_mode": "allow_fp32_to_fp16", "log": "debug", "op_debug_level": self.op_debug_level}

        return_code = self.run_atc(om_params)
        return return_code


atc = AtcOm()
